# Multi-agent Environment Visualization

from metadrive.envs.marl_envs import MultiAgentRoundaboutEnv, MultiAgentBottleneckEnv, MultiAgentIntersectionEnv, MultiAgentParkingLotEnv, MultiAgentTollgateEnv
from metadrive.examples import expert

env_classes = [MultiAgentRoundaboutEnv, 
        MultiAgentBottleneckEnv, 
        MultiAgentIntersectionEnv, 
        MultiAgentParkingLotEnv, 
        MultiAgentTollgateEnv]

frames = []

for env_class in env_classes:
    env = env_class()
    print(f"Starting the environment {env}.")
    env.reset()

    tm={"__all__":False}

    for i in range(100):
        if tm["__all__"]:
            frames.append(frame)
            continue

        action = env.action_space.sample()
        for a in action.values():
            a[-1] = 1.0
        o,r,tm,tc,i = env.step(action)
        frame = env.render(mode='top_down',
                           scaling=4,
                           camera_position=env.current_map.get_center_point(),
                           screen_size=(500, 500)
        )
        frames.append(frame)

    env.close()

# render image
print("\nGenerate gif...")
import pygame
import numpy as np
from PIL import Image
import os

HOME = os.path.dirname(os.path.realpath(__file__))

images = [frame for frame in frames]
images = [Image.fromarray(image) for image in images]
images[0].save(f"{HOME}/demo.gif", save_all=True, append_images=images[1:], loop=0)

